#pragma once
#ifndef NORM_ACTV_H_
#define NORM_ACTV_H_

#include "stream_tools.h"

template<
	unsigned N_IO,		// # of IO
	unsigned N_CH,		// # of CHannel
	unsigned R_SHIFT,	// nBit of Right Shift
	unsigned VEC_LEN	// ROW * COL
>
void norm_relu(
	data_stream<N_IO * BIT_CONV>& in, 
	data_stream<N_IO * BIT_ACTV>& out, 
	const ap_uint<N_IO * BIT_BIAS> bias[N_CH / N_IO],
	const ap_uint<N_IO * BIT_MULT> mult[N_CH / N_IO]
)
{
	static_assert(N_CH >= N_IO, "norm_actv");
	static_assert(N_CH % N_IO == 0, "norm_actv");
	constexpr unsigned FOLD = N_CH / N_IO;
	constexpr unsigned ITERS = VEC_LEN * FOLD;
	constexpr unsigned MAXOUT = (1 << BIT_ACTV) - 1;

	assert(in.size() == VEC_LEN * FOLD);
	assert(out.empty());

	#pragma HLS DATAFLOW

	unsigned f = 0;
	for (unsigned it = 0; it < ITERS; ++it)
	{
		#pragma HLS PIPELINE II=1
		ap_uint<N_IO * BIT_CONV> in_buf = in.read();
		ap_uint<N_IO * BIT_BIAS> b_buf = bias[f];
		ap_uint<N_IO * BIT_MULT> m_buf = mult[f];
		ap_uint<N_IO * BIT_ACTV> out_buf;

		for (unsigned i = 0; i < N_IO; ++i)
		{
			#pragma HLS UNROLL
			// UnRoll
			ap_int<BIT_CONV> a = in_buf(SLICE(BIT_CONV, i));
			ap_int<BIT_BIAS> b = b_buf(SLICE(BIT_BIAS , i));
			ap_int<BIT_MULT> m = m_buf(SLICE(BIT_MULT , i));
			ap_int<BIT_BIAS + BIT_MULT + 1> x = (a + b) * m >> R_SHIFT;
			ap_int<BIT_BIAS + BIT_MULT + 1> x_shift;
			if(x.range(0,0) > 0)
			{
				x_shift = x >> 1;
				x_shift = x_shift + 1;
			}
			else
			{
				x_shift = x >> 1;
			}
			ap_uint<BIT_ACTV> y = 0;
			if (x_shift > 0)
			{
				y = x_shift < MAXOUT ? ap_uint<BIT_ACTV>(x_shift) : ap_uint<BIT_ACTV>(MAXOUT);
			}
			else
			{
				y = 0;
			}
			out_buf(SLICE(BIT_ACTV, i)) = y;
		}
		out.write(out_buf);
		f = (f != FOLD - 1) ? f + 1 : 0;
	}
	assert(in.empty());
	assert(out.size() == VEC_LEN * FOLD);
	return;
}

template<
	unsigned N_IO,		// # of IO
	unsigned N_CH,		// # of CHannel
	unsigned R_SHIFT,	// nBit of Right Shift
	unsigned VEC_LEN	// ROW * COL
>
void norm(
	data_stream<N_IO * BIT_CONV>& in, 
	data_stream<N_IO * BIT_ACTV>& out, 
	const ap_uint<N_IO * BIT_BIAS> bias[N_CH / N_IO],
	const ap_uint<N_IO * BIT_MULT> mult[N_CH / N_IO]
)
{
	static_assert(N_CH >= N_IO, "norm_actv");
	static_assert(N_CH % N_IO == 0, "norm_actv");
	constexpr unsigned FOLD = N_CH / N_IO;
	constexpr unsigned ITERS = VEC_LEN * FOLD;
	constexpr signed MINOUT = - (1 << (BIT_ACTV-1));
	constexpr signed MAXOUT = (1 << (BIT_ACTV-1)) - 1;

	assert(in.size() == VEC_LEN * FOLD);
	assert(out.empty());

	#pragma HLS DATAFLOW
	unsigned pixel_i = 0;
	unsigned f = 0;
	for (unsigned it = 0; it < ITERS; ++it)
	{
		#pragma HLS PIPELINE II=1
		ap_uint<N_IO * BIT_CONV> in_buf = in.read();
		ap_uint<N_IO * BIT_BIAS> b_buf = bias[f];
		ap_uint<N_IO * BIT_MULT> m_buf = mult[f];
		ap_uint<N_IO * BIT_ACTV> out_buf;
		for (unsigned i = 0; i < N_IO; ++i)
		{
			#pragma HLS UNROLL
			// UnRoll
			ap_int<BIT_CONV> a = in_buf(SLICE(BIT_CONV, i));
			ap_int<BIT_BIAS> b = b_buf(SLICE(BIT_BIAS , i));
			ap_int<BIT_MULT> m = m_buf(SLICE(BIT_MULT , i));
			ap_int<BIT_BIAS + BIT_MULT + 1> x = (a + b) * m >> R_SHIFT;
			ap_int<BIT_BIAS + BIT_MULT + 1> x_shift;
			if(x.range(0,0) > 0)
			{
				x_shift = x >> 1;
				x_shift = x_shift + 1;
			}
			else
			{
				x_shift = x >> 1;
			}
			ap_int<BIT_ACTV> y = 0;
			if (x_shift > MAXOUT)
			{
				y = ap_int<BIT_ACTV>(MAXOUT);
			}
			else if(x_shift < MINOUT)
			{
				y = ap_int<BIT_ACTV>(MINOUT);
			}
			else
			{
				y = ap_int<BIT_ACTV>(x_shift);
			}
			out_buf(SLICE(BIT_ACTV, i)) = y;

		}
		out.write(out_buf);
		f = (f != FOLD - 1) ? f + 1 : 0;
	}
	assert(in.empty());
	assert(out.size() == VEC_LEN * FOLD);
	return;
}

#endif
